use crate::{HadesCryptoError, HadesCryptoResult};
use crypto::digest::Digest;
use crypto::sha1::Sha1;
use crypto::sha2::Sha256;
/// 利用标准的公钥和私钥进行加签、验签、加密、解密
/// PKS#1 就是带有-----BEGIN RSA PRIVATE KEY-----的格式
/// PKS#8 就是 带有-----BEGIN PRIVATE KEY-----的格式
/// 公钥默认都是PKCS8格式的。
use rsa::{Hash, PaddingScheme, PublicKey, RSAPrivateKey, RSAPublicKey};
use std::iter::repeat;

pub enum HashType {
    Sha1,
    Sha256,
}

pub enum KeyType {
    PKS1,
    PKS8,
}

pub fn format_key(key: &str) -> String {
    key.lines()
        .filter(|line| !line.starts_with("-"))
        .fold(String::new(), |mut data, line| {
            data.push_str(&line);
            data
        })
}

/// 加签
pub fn rsa_sign(
    content: &str,
    private_key: &str,
    key_type: KeyType,
    hash_type: HashType,
) -> HadesCryptoResult<String> {
    let der_bytes = base64::decode(format_key(private_key))?;
    let private_key = match key_type {
        KeyType::PKS1 => RSAPrivateKey::from_pkcs1(&der_bytes)
            .map_err(|e| HadesCryptoError::PrivateKeyDecode(e.to_string()))?,
        KeyType::PKS8 => RSAPrivateKey::from_pkcs8(&der_bytes)
            .map_err(|e| HadesCryptoError::PrivateKeyDecode(e.to_string()))?,
    };

    let (hash, buf) = match hash_type {
        HashType::Sha1 => {
            let hash = Hash::SHA1;
            let mut hasher = Sha1::new();
            hasher.input_str(content);
            let mut buf = repeat(0)
                .take((hasher.output_bits() + 7) / 8)
                .collect::<Vec<u8>>();
            hasher.result(&mut buf);
            (hash, buf)
        }

        HashType::Sha256 => {
            let hash = Hash::SHA2_256;
            let mut hasher = Sha256::new();
            hasher.input_str(content);
            let mut buf = repeat(0)
                .take((hasher.output_bits() + 7) / 8)
                .collect::<Vec<u8>>();
            hasher.result(&mut buf);
            (hash, buf)
        }
    };

    private_key
        .sign(PaddingScheme::PKCS1v15Sign { hash: Some(hash) }, &buf)
        .map_err(|e| HadesCryptoError::Sign(e.to_string()))
        .and_then(|ret| Ok(base64::encode(ret)))
}

/// 验签
pub fn rsa_verify(
    content: &str,
    public_key: &str,
    sign: &str,
    hash_type: HashType,
) -> HadesCryptoResult<()> {
    let public_key = base64::decode(format_key(public_key))?;
    let sign = base64::decode(sign)?;
    let public_key = RSAPublicKey::from_pkcs8(&public_key)
        .map_err(|e| HadesCryptoError::PublicKeyDeCode(e.to_string()))?;

    // 进行摘要
    let hash;
    let mut buf;
    match hash_type {
        HashType::Sha1 => {
            hash = Hash::SHA1;
            let mut hasher = Sha1::new();
            hasher.input_str(content);
            buf = repeat(0)
                .take((hasher.output_bits() + 7) / 8)
                .collect::<Vec<u8>>();
            hasher.result(&mut buf);
        }
        HashType::Sha256 => {
            hash = Hash::SHA2_256;
            let mut hasher = Sha256::new();
            hasher.input_str(content);
            buf = repeat(0)
                .take((hasher.output_bits() + 7) / 8)
                .collect::<Vec<u8>>();
            hasher.result(&mut buf);
        }
    }

    public_key
        .verify(
            PaddingScheme::PKCS1v15Sign { hash: Some(hash) },
            &buf,
            &sign,
        )
        .map_err(|e| HadesCryptoError::Verify(e.to_string()))
}

/// 加密,使用公钥加密
pub fn rsa_encrypt(content: &str, public_key: &str) -> HadesCryptoResult<String> {
    let public_key = base64::decode(format_key(public_key))?;
    let public_key = RSAPublicKey::from_pkcs8(&public_key)
        .map_err(|e| HadesCryptoError::PublicKeyDeCode(e.to_string()))?;
    public_key
        .encrypt(
            &mut rand::thread_rng(),
            PaddingScheme::PKCS1v15Encrypt,
            content.as_bytes(),
        )
        .map_err(|e| HadesCryptoError::Encrypt(e.to_string()))
        .and_then(|res| Ok(base64::encode(res)))
}

/// 解密
pub fn rsa_decrypt(
    pks_type: KeyType,
    private_key: &str,
    encoded: &str,
) -> HadesCryptoResult<String> {
    let private_key = base64::decode(format_key(private_key))?;
    let private_key = match pks_type {
        KeyType::PKS1 => RSAPrivateKey::from_pkcs1(&private_key)
            .map_err(|e| HadesCryptoError::PrivateKeyDecode(e.to_string()))?,
        KeyType::PKS8 => RSAPrivateKey::from_pkcs8(&private_key)
            .map_err(|e| HadesCryptoError::PrivateKeyDecode(e.to_string()))?,
    };

    let ori_data = base64::decode(encoded)?;
    private_key
        .decrypt(PaddingScheme::PKCS1v15Encrypt, &ori_data)
        .map_err(|e| HadesCryptoError::Decrypt(e.to_string()))
        .and_then(|ori_data| {
            String::from_utf8(ori_data).map_err(|e| HadesCryptoError::Other(e.to_string()))
        })
}

#[cfg(test)]
mod tests {
    use crate::crypto::{rsa_decrypt, rsa_encrypt, rsa_sign, rsa_verify, HashType, KeyType};
    use std::fs;

    #[test]
    fn test_sign() {
        let content = "hello world women 世界你好";
        let private_pkcs1 = fs::read_to_string("./pri_pkcs1.pem").unwrap();
        let public_pkcs1 = fs::read_to_string("./pub_pkcs1.pem").unwrap();
        let sign = rsa_sign(content, &private_pkcs1, KeyType::PKS8, HashType::Sha1);
        if let Err(e) = sign {
            println!("e: {}", e);
            return;
        }
        println!("rsa_sign: {:?}", sign);
        let ret = rsa_verify(content, &public_pkcs1, &sign.unwrap(), HashType::Sha1);
        if let Err(e) = ret {
            println!("rsa verify failed: {}", e);
            return;
        } else {
            println!("rsa verify success");
        }

        let encoded_data = rsa_encrypt(&content, &public_pkcs1);
        if let Err(e) = encoded_data {
            println!("e: {}", e);
            return;
        }
        println!("encoded_data: {:?}", encoded_data);
        let ori_data = rsa_decrypt(KeyType::PKS1, &private_pkcs1, &encoded_data.unwrap());
        if let Err(e) = ori_data {
            println!("e: {}", e);
            return;
        }
        println!("oridata: {:?}", ori_data);
    }
}
